#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdio.h>

#define KEY_LENGTH 256  // Key table length


// Generate key table
void generate_key_table(unsigned char *key_table) {
    const char *special_chars = "!@#$%^&*()_+-=[]{}|;:,.<>?`~";
    int special_chars_len = strlen(special_chars);
    
    // 使用特殊字符循环填充密码表
    for (int i = 0; i < KEY_LENGTH; i++) {
        key_table[i] = special_chars[i % special_chars_len];
    }
}

// 保存密码表到文件
void save_key_table(const char *filename, unsigned char *key_table) {
    FILE *fp = fopen(filename, "wb");
    if (fp == NULL) {
        printf("Failed to create key table file\n");
        return;
    }
    fwrite(key_table, 1, KEY_LENGTH, fp);
    fclose(fp);
}

// 从文件加载密码表
int load_key_table(const char *filename, unsigned char *key_table) {
    FILE *fp = fopen(filename, "rb");
    if (fp == NULL) {
        printf("Failed to open key table file\n");
        return 0;
    }
    size_t read_size = fread(key_table, 1, KEY_LENGTH, fp);
    fclose(fp);
    return read_size == KEY_LENGTH;
}

// 加密/解密文件
void crypto_file(const char *input_file, const char *output_file, unsigned char *key_table) {
    FILE *fp_in = fopen(input_file, "rb");
    if (fp_in == NULL) {
        printf("Failed to open input file\n");
        return;
    }

    FILE *fp_out = fopen(output_file, "wb");
    if (fp_out == NULL) {
        printf("Failed to create output file\n");
        fclose(fp_in);
        return;
    }

    // 使用更大的缓冲区以提高性能
    unsigned char buffer[8192];  // 8KB 缓冲区
    int key_index = 0;
    size_t bytes_read;
    
    while ((bytes_read = fread(buffer, 1, sizeof(buffer), fp_in)) > 0) {
        // 对缓冲区中的每个字节进行加密/解密
        for (size_t i = 0; i < bytes_read; i++) {
            buffer[i] ^= key_table[key_index];
            key_index = (key_index + 1) % KEY_LENGTH;
        }
        
        fwrite(buffer, 1, bytes_read, fp_out);
    }

    fclose(fp_in);
    fclose(fp_out);
    printf("complete:100%%\n");  // 添加完成提示
}

// ... existing code ...

int main() {
    unsigned char key_table[KEY_LENGTH];
    char choice;
    char input_file[256];
    char output_file[256];

    printf("Please select operation:\n");
    printf("1. Encrypt file\n");
    printf("2. Decrypt file\n");
    if (scanf(" %c", &choice) != 1) {
        printf("Invalid input. Please try again.\n");
        return 1;
    }

    switch (choice) {
        case '0':
            // Generate new key table
            generate_key_table(key_table);
            save_key_table("key_table.bin", key_table);
            printf("Key table has been generated and saved to key_table.bin\n");
            break;

        case '1':
        case '2':
            // 加载密码表

            generate_key_table(key_table);
            save_key_table("key_table.bin", key_table);

            if (!load_key_table("key_table.bin", key_table)) {
                printf("Error!!\n");
                return 1;
            }

            printf("Source file name:");
            if (scanf("%s", input_file) != 1) {
                printf("Invalid input. Please try again.\n");
                return 1;
            }
            printf("Destination file name:");
            if (scanf("%s", output_file) != 1) {
                printf("Invalid input. Please try again.\n");
                return 1;
            }

            crypto_file(input_file, output_file, key_table);
            remove("key_table.bin");
            printf("Complete\n");
            break;

        default:
            printf("Error\n");
            return 1;
    }

    return 0;
}